from fastapi import APIRouter, WebSocket
from pydantic import BaseModel
from typing import Optional, Dict
import uuid
from sqllite_conn import cursor, conn  # 继承现有数据库连接
from primisai.nexus.config import load_yaml_config, AgentFactory

router = APIRouter(prefix="/api/task", tags=["Task Management"])

class TaskCreateRequest(BaseModel):
    """任务创建请求模型
    Attributes:
        task: 任务描述
        subtasks: 子任务列表（可选）
    """
    task: str
    subtasks: Optional[list[str]] = None

class TaskQueryRequest(BaseModel):
    """任务查询请求模型
    Attributes:
        status: 任务状态过滤条件
    """
    status: str

class StartTaskRequest(BaseModel):
    """任务启动请求模型
    Attributes:
        type: 任务类型
        task_id: 父任务ID
        task: 任务描述
        subtask_list: 子任务字典（子任务ID:子任务描述）
    """
    type: str
    task_id: str
    task: str
    subtask_list: Dict[str, str]

class TaskUpdateRequest(BaseModel):
    """任务更新请求模型
    Attributes:
        task_id: 要更新的任务ID
        new_task: 更新后的任务描述
    """
    task_id: str
    new_task: str

class TaskDeleteRequest(BaseModel):
    """任务删除请求模型
    Attributes:
        task_id: 要删除的任务ID
    """
    task_id: str

class TaskStopRequest(BaseModel):
    """任务停止请求模型
    Attributes:
        task_id: 要停止的任务ID
    """
    type: str
    task_id: str

@router.get("/query")
async def query_task(status: str):
    """
    查询任务接口
    - status: 过滤任务状态参数
    - 返回父子任务关联结构
    """
    cursor.execute("SELECT * FROM tasks WHERE status=? AND father_node IS NULL", (status,))
    parent_tasks = cursor.fetchall()
    
    result = []
    for task in parent_tasks:
        cursor.execute("SELECT * FROM tasks WHERE father_node=?", (task[1],))
        subtasks = cursor.fetchall()
        
        result.append({
            "task_id": task[1],
            "task": task[2],
            "status": task[3],
            "subtask_list": [{
                "subtask_id": st[1],
                "subtask": st[2],
                "status": st[3],
                "result": st[4]
            } for st in subtasks]
        })
    
    return {"code": 200, "message": "success", "data": result}



@router.post("/new")
async def create_task(request: TaskCreateRequest):
    """
    创建新任务接口
    - 生成UUID任务ID
    - 初始化任务状态为pending
    - 调用Nexus框架进行智能任务拆分
    """
    task_id = str(uuid.uuid4())
    cursor.execute("INSERT INTO tasks (task_id, task, status) VALUES (?, ?, 'pending')",
                   (task_id, request.task))
    conn.commit()

    # Nexus任务拆分（需确保agent_config/task_config.yaml已配置）
    try:
        config = load_yaml_config('agent_config/task_config.yaml')
        agent = AgentFactory.create_from_config(config)
        # agent确定有breakdown_task方法？
        subtask_list = agent.breakdown_task(request.task)

        # 生成子任务并建立关联
        for subtask in subtask_list:
            subtask_id = str(uuid.uuid4())
            cursor.execute("INSERT INTO tasks (task_id, task, status, father_node) VALUES (?, ?, 'pending', ?)",
                          (subtask_id, subtask, task_id))

        conn.commit()

    except Exception as e:
        print(f"Nexus任务拆分失败: {str(e)}")

    return {"code": 200, "message": "success", "data": {"task_id": task_id}}

@router.post("/start")
async def start_task(request: StartTaskRequest):
    """
    启动任务接口
    - 参数校验：task_id必须存在
    - 更新父任务状态为running
    - 批量插入子任务记录
    """
    cursor.execute("UPDATE tasks SET status='running' WHERE task_id=?",
                   (request.task_id,))
    
    for subtask_id, subtask in request.subtask_list.items():
        cursor.execute("INSERT INTO tasks (task_id, task, status, father_node) VALUES (?, ?, 'pending', ?)",
                       (subtask_id, subtask, request.task_id))
    conn.commit()
    return {"code": 200, "message": "success"}


@router.post("/update")
async def update_task(request: TaskUpdateRequest):
    """
    更新任务接口
    - 参数校验：task_id必须存在
    - 更新任务描述
    - 事务提交保证数据一致性
    """
    cursor.execute("UPDATE tasks SET task=? WHERE task_id=?",
                   (request.new_task, request.task_id))
    conn.commit()
    return {"code": 200, "message": "success"}

@router.post("/delete")
async def delete_task(request: TaskDeleteRequest):
    """
    删除任务接口
    - 级联删除父任务及关联子任务
    - 使用father_node查询确保删除所有关联记录
    - 事务提交保证数据一致性
    """
    cursor.execute("DELETE FROM tasks WHERE task_id=?", (request.task_id,))
    cursor.execute("DELETE FROM tasks WHERE father_node=?", (request.task_id,))
    conn.commit()
    return {"code": 200, "message": "success"}



@router.post("/stop")
async def stop_task(request: TaskStopRequest):
    """
    停止任务接口
    - 参数校验：task_id必须存在
    - 更新任务状态为stopped
    - 级联更新关联子任务状态
    """
    cursor.execute("UPDATE tasks SET status='stopped' WHERE task_id=?", (request.task_id,))
    cursor.execute("UPDATE tasks SET status='stopped' WHERE father_node=?", (request.task_id,))
    conn.commit()
    return {"code": 200, "message": "success"}